#!/opt/local/bin/python

# Copyright Stephen Molloy, 2012
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import pyparsing as pyp
import csv
from scipy import nan, isnan

class DSKParser(object):
    """Instances of this class contain all the data extracted from a DSK
    file, and contain all the information necessary to specify and simulate
    a klystron using the AJDisk engine.
    The user will rarely (if at all) have to use any of the methods contained
    in this class, and will normally only need to pass instances to AJDiskAnalysis
    classes.
    
    parserobj = DSKParser(fname)
    fname is the filename of a correctly formatted DSK file."""
    def __init__(self, fname):
        self.fname = fname
        self.doParsing()
        self.makeOPs()
        self.makeCavs()
        self.makeSim()

    def doParsing(self):
        COLON = pyp.Suppress(":")
        title = pyp.Word(pyp.alphanums+"_") + COLON
        line = title.setResultsName("title") + pyp.restOfLine.setResultsName("var")
        
        self.dskInput = {}
        filedata = open(self.fname).read()
        for t,s,e in line.scanString(filedata):
            t.var = t.var.strip()
            self.dskInput[t.title[0]] = t.var.split()

        if not "r_over_q" in self.dskInput:
            self.dskInput['r_over_q']   = []
            self.dskInput['m_input']    = []
            self.dskInput['qe']         = []
            self.dskInput['q0']         = []
            self.dskInput['tuned_f']    = []
            self.dskInput['gap_z']      = []
            self.dskInput['gap_length'] = []
            self.dskInput['type']       = []
            foundcavdat = False
            file = open(self.fname)
            for line in csv.reader(file, delimiter=" "):
                newline = " ".join(line).strip().split()
                try:
                    title = newline[0]
                except IndexError:
                    title = nan
                if title == "Cavity":
                    break
            for line in csv.reader(file, delimiter=" "):
                newline = " ".join(line).strip().split()
                cavnum = int(newline[0]) - 1
                self.dskInput['r_over_q'].append( float(newline[2]) )
                self.dskInput['m_input'].append( float(newline[8]) )
                self.dskInput['qe'].append( float(newline[3]) )
                self.dskInput['q0'].append( float(newline[4]) )
                self.dskInput['tuned_f'].append( float(newline[5]) )
                self.dskInput['gap_z'].append( float(newline[6]) )
                self.dskInput['gap_length'].append( float(newline[7]) )
                self.dskInput['type'].append( float(newline[1]) )

    def makeOPs(self):
        V0  = float(self.dskInput['v0'][0])
        I0  = float(self.dskInput['i0'][0])
        f   = float(self.dskInput['f'][0])
        a   = float(self.dskInput['r_drift'][0])
        b   = float(self.dskInput['r_beam'][0])
        Pin = float(self.dskInput['pin'][0])
        self.opParam = operatingParams(V0=V0, I0=I0, F0=f, Pin=Pin, a=a, b=b)

    def makeCavs(self):
        self.cavs = []
        for cavnum in xrange(int( self.dskInput['cavities'][0])):
            self.cavs.append(DSKcavity(
                RoverQ   = float( self.dskInput['r_over_q'][cavnum]),
                M        = float( self.dskInput['m_input'][cavnum]),
                Qe       = float( self.dskInput['qe'][cavnum]),
                Qo       = float( self.dskInput['q0'][cavnum]),
                f        = float( self.dskInput['tuned_f'][cavnum]),
                z        = float( self.dskInput['gap_z'][cavnum]),
                d        = float( self.dskInput['gap_length'][cavnum]),
                harmonic = int( self.dskInput['type'][cavnum])))

    def makeSim(self):
        self.simobj = DSKsim(
                disks   = int(self.dskInput['disks'][0]),
                rings   = int(self.dskInput['rings'][0]),
                steps   = int(self.dskInput['steps'][0]),
                maxiter = int(self.dskInput['max_iter'][0]))

class operatingParams(object):
    """Instances of this class are used by DSKParser to hold all the information relating
    to the operating parameters of the klystron to be simulated using AJDisk.
    This class is used by DSKParser to hold information, and it is unlikely that the user
    will ever need access to this class outside of that context."""
    def __init__(self, V0, I0, F0, Pin, a=nan, b=nan, xb=nan, yb=nan, xd=nan, yd=nan):
        self.V0  = V0
        self.I0  = I0
        self.F0  = F0
        self.Pin = Pin
        self.a   = a
        self.b   = b
        self.xb  = xb
        self.yb  = yb
        self.xd  = xd
        self.yd  = yd
        self.perv = I0 / (V0**1.5)

class DSKcavity(object):
    """Instances of this class are used by DSKParser objects to hold information on
    individual cavities within a klystron.
    It is unlikely that the user will ever need access to this class."""
    def __init__(self, RoverQ, M, Qe, Qo, f, z, d, harmonic):
        self.RoverQ   = RoverQ
        self.M        = M
        self.Qe       = Qe
        self.Qo       = Qo
        self.f        = f
        self.z        = z
        self.d        = d
        self.harmonic = harmonic

class DSKsim(object):
    """Instances of this class are used by DSKParser to hold information relevant to
    the large signal simulation of klystrons.
    It is unlikely that users will ever need direct access to this class."""
    def __init__(self, disks, rings, steps, maxiter):
        self.disks   = disks
        self.rings   = rings
        self.steps   = steps
        self.maxiter = maxiter

